#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import re
import argparse
from PyPDF2 import PdfReader
import logging

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

def extract_headings_from_pdf(pdf_path):
    """从PDF文件中提取可能的标题"""
    try:
        logger.info(f"正在处理PDF文件: {pdf_path}")
        reader = PdfReader(pdf_path)
        headings = []
        
        # 尝试从PDF的书签中提取目录结构
        if reader.outline:
            logger.info(f"找到PDF内置书签")
            
            def process_outline_item(item, level=0):
                result = []
                if isinstance(item, dict):
                    # 处理单个书签项
                    if '/Title' in item:
                        title = item['/Title']
                        page_num = None
                        if '/Page' in item:
                            page_num = reader.get_destination_page_number(item) + 1  # 页码从0开始，需要+1
                        elif '/Dest' in item and item['/Dest']:
                            try:
                                page_ref = item['/Dest'][0]
                                for i, page in enumerate(reader.pages):
                                    if page.get_object() == page_ref:
                                        page_num = i + 1
                                        break
                            except:
                                pass
                        
                        if page_num is not None:
                            result.append((title, page_num, level))
                        else:
                            result.append((title, 1, level))  # 如果无法获取页码，默认为第1页
                        
                    # 处理子书签
                    if '/First' in item:
                        child = item['/First']
                        while child:
                            result.extend(process_outline_item(child, level + 1))
                            if '/Next' in child:
                                child = child['/Next']
                            else:
                                break
                
                elif isinstance(item, list):
                    # 处理书签列表
                    for subitem in item:
                        result.extend(process_outline_item(subitem, level))
                
                return result
            
            # 处理整个书签树
            if isinstance(reader.outline, list):
                for item in reader.outline:
                    headings.extend(process_outline_item(item))
            else:
                headings.extend(process_outline_item(reader.outline))
            
            # 调试信息
            logger.info(f"提取到 {len(headings)} 个书签")
            
        else:
            logger.info(f"PDF没有内置书签，尝试从文本内容中提取标题")
            # 尝试从文本内容中提取标题
            for page_num, page in enumerate(reader.pages, 1):
                text = page.extract_text()
                if text:
                    # 尝试查找可能的标题（简单示例，可以根据实际情况调整）
                    lines = text.split('\n')
                    for line in lines:
                        # 假设标题是较短的行，并且可能以数字或第、章等开头
                        line = line.strip()
                        if 5 <= len(line) <= 100 and (re.match(r'^[第章节\d]+', line) or re.match(r'^\d+[\.\s]', line)):
                            headings.append((line, page_num, 0))
        
        return headings
    except Exception as e:
        logger.error(f"处理文件 {pdf_path} 时出错: {str(e)}")
        logger.exception(e)
        return []

def generate_toc(directory_path, output_file='toc.txt', recursive=False):
    """生成指定目录下所有PDF文件的目录"""
    logger.info(f"开始生成目录，路径: {directory_path}")
    
    if not os.path.exists(directory_path):
        logger.error(f"目录不存在: {directory_path}")
        return
    
    all_pdfs = []
    
    # 收集所有PDF文件
    if recursive:
        for root, _, files in os.walk(directory_path):
            for file in files:
                if file.lower().endswith('.pdf'):
                    all_pdfs.append(os.path.join(root, file))
    else:
        all_pdfs = [os.path.join(directory_path, f) for f in os.listdir(directory_path) 
                   if os.path.isfile(os.path.join(directory_path, f)) and f.lower().endswith('.pdf')]
    
    all_pdfs.sort()  # 按文件名排序
    
    logger.info(f"找到 {len(all_pdfs)} 个PDF文件")
    
    # 生成目录
    with open(output_file, 'w', encoding='utf-8') as f:
        f.write("# PDF文档目录\n\n")
        
        for pdf_path in all_pdfs:
            rel_path = os.path.relpath(pdf_path, directory_path)
            f.write(f"## {rel_path}\n\n")
            
            headings = extract_headings_from_pdf(pdf_path)
            
            if headings:
                for title, page_num, level in headings:
                    indent = "  " * level
                    f.write(f"{indent}- {title} (第{page_num}页)\n")
            else:
                f.write("未能提取到目录信息\n")
            
            f.write("\n---\n\n")
    
    logger.info(f"目录已生成: {output_file}")

def main():
    parser = argparse.ArgumentParser(description='生成PDF文件的目录')
    parser.add_argument('directory', help='PDF文件所在的目录路径')
    parser.add_argument('-o', '--output', default='toc.txt', help='输出文件名 (默认: toc.txt)')
    parser.add_argument('-r', '--recursive', action='store_true', help='是否递归搜索子目录')
    
    args = parser.parse_args()
    
    generate_toc(args.directory, args.output, args.recursive)

if __name__ == "__main__":
    main() 